#!/bin/bash

declare -A epochs=(["mrpc"]=30 ["qnli"]=25 ["rte"]=80 ["sst2"]=60 ["stsb"]=40 ["cola"]=80)
declare -A bs=(["mrpc"]=64 ["qnli"]=64 ["rte"]=64 ["sst2"]=64 ["stsb"]=64 ["cola"]=64)
declare -A ml=(["mrpc"]=256 ["qnli"]=256 ["rte"]=512 ["sst2"]=256 ["stsb"]=256 ["cola"]=256)
declare -A lr=(["mrpc"]="4e-4" ["qnli"]="4e-4" ["rte"]="4e-4" ["sst2"]="5e-4" ["stsb"]="4e-4" ["cola"]="4e-4")
declare -A metrics=(["mrpc"]="accuracy" ["qnli"]="accuracy" ["rte"]="accuracy" ["sst2"]="accuracy" ["stsb"]="pearson" ["cola"]="matthews_correlation")
declare -A target=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
declare -A eval_steps_list=(["Ada"]=100 ["Mask"]=100)
#declare -A reg_orth_coef=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A init_warmup=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A final_warmup=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A mask_interval=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A lora_alpha=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A max_seq_length=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A learning_rate=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A num_train_epochs=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A warmup_steps=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A cls_dropout=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
#declare -A weight_decay=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")

export WANDB_MODE=offline

run(){
  eval_steps=${eval_steps_list[$3]}
  learning_rate=${lr[$1]}
  num_train_epochs=${epochs[$1]}
  per_device_train_batch_size=${bs[$1]}
  gradient_accumulation_steps=1
  task_name=$1
  rank=$2
  mode=$3
  seed=$4
  Iidea=$5
  RS_flag=$6
  lora_alpha="16"
  target_modules=${target[$4]}
  lora_dropout=0.05
  lora_bias=none
  lora_task_type=SEQ_CLS
  wandb_project=msplora-reproduce
  share=false
  is_train_in_stage=false
  wandb_run_name=roberta-lora-${mode}-${task_name}-stage-${is_train_in_stage}-r-${rank}-n-alpha-16-seed-${seed}-bs-${per_device_train_batch_size}-lr-${learning_rate}-epochs-${num_train_epochs}

  exp_dir=./roberta_glue_reproduce_${mode}/${wandb_run_name}
  new_dir=./output/glue/${mode}-${Iidea}-Rs${RS_flag}_0.0001_L2/${task_name}-seed${seed}-rank${rank}
  False_new_dir=./output/glue/${mode}-${Iidea}-Rs${RS_flag}/${task_name}-seed${seed}-rank${rank}


  python examples/text-classification/run_glue.py \
  --model_name_or_path microsoft/deberta-v3-base \
  --task_name ${task_name} \
  --apply_adalora --apply_lora --lora_type svd \
  --RS_flag ${RS_flag} \
  --target_rank 2   --lora_r 4   \
  --reg_orth_coef 0.1 \
  --init_warmup 600 --final_warmup 1800 --mask_interval 1 \
  --beta1 0.85 --beta2 0.85 \
  --mode ${mode} --Iidea ${Iidea} \
  --lora_module query,key,value,intermediate,layer.output,attention.output \
  --lora_alpha 32 \
  --do_train --do_eval --max_seq_length 320 \
  --per_device_train_batch_size 32 --learning_rate 1e-3 \
  --num_train_epochs 30 --warmup_ratio 0.1 \
  --cls_dropout 0.0 --weight_decay 0.01 \
  --evaluation_strategy steps --eval_steps 300 \
  --save_strategy steps --save_steps 3000 \
  --logging_steps 100 \
  --report_to tensorboard \
  --seed ${seed} \
  --root_output_dir ${new_dir} \
  --overwrite_output_dir
}
#task_base=('mrpc' 'qnli' 'rte' 'sst2' 'stsb' 'cola')
task_base=('mrpc')
seeds=(40)
#init_rank=(4 6 12)
#end_rank=(2 4 8)
ranks=(4)
modes=("Ada")
Iideas=("none")
RS=(True)


for task in "${task_base[@]}"; do
  for seed in "${seeds[@]}"; do
    for rank in "${ranks[@]}"; do
      for RS_flag in "${RS[@]}"; do
        for mode in "${modes[@]}"; do
          if [ "$mode" = "Ada" ]; then
            Iidea="none"
            echo "Running with task: $task, seed: $seed, rank: $rank, mode: $mode, Iidea: $Iidea, RS_flag: $RS_flag"
            run "$task" "$rank" "$mode" "$seed" "$Iidea" "$RS_flag"
          else
            for Iidea in "${Iideas[@]}"; do
              echo "Running with task: $task, seed: $seed, rank: $rank, mode: $mode, Iidea: $Iidea, RS_flag: $RS_flag"
              run "$task" "$rank" "$mode" "$seed" "$Iidea" "$RS_flag"
            done
          fi
        done
      done
    done
  done
done
